import sys
import os
import cPickle as pickle
import numpy as np
import load
load = reload( load )
TOL = 0.001
def estimate_P1( plcurve ):
    [(x1,y1),(x2,y2),(x3,y3)] = plcurve[:3]
    [(x4,y4),(x5,y5),(x6,y6)] = plcurve[-3:]
    c1 = np.poly1d( np.polyfit( [x1,x2,x3],[y1,y2,y3],1) )
    c2 = np.poly1d( np.polyfit( [x4,x5,x6],[y4,y5,y6],1) )
    m1 = np.array( [[c1.coeffs[0], -1], [c2.coeffs[0], -1]] )
    m2 = np.array( [[-c1.coeffs[1]], [-c2.coeffs[1]]] )
    p1 = float( np.linalg.solve( m1, m2 )[0] )
    return p1

class Study(object):

    def __init__( self ):
        self.name = ''
        self.tmp_dir =  r'C:\Temp\abaqus'
        self.source_python = r'\\MA-CASTRO\abaqus-conecyl-python'
        self.ccs = [] # ccs --> conecyls - conecylinders
        self.rannames = []
        self.xyplots = []

        self.group = 0
        self.study_dir  = ''
        self.python_dir = ''
        self.output_dir = ''
        self.run_file_name = ''
        self.session = None
        self.p1 = None

    def rebuild( self ):
        self.study_dir  = os.path.join( self.tmp_dir  , self.name )
        self.python_dir = os.path.join(\
            self.study_dir,  'abaqus-conecyl-python' )
        self.output_dir = os.path.join( self.study_dir, 'outputs' )
        for cc in self.ccs:
            cc.tmp_dir = self.tmp_dir
            cc.study = self
            cc.study_dir = self.study_dir
            cc.output_dir = self.output_dir
            cc.rebuild()

    def reset_folders( self ):
        pass

    def configure_folders( self ):
        self.rebuild()
        print 'configuring folders...'
        print '\t' + self.python_dir
        print '\t' + self.output_dir
        if not os.path.isdir( self.output_dir ):
            os.makedirs( self.output_dir )
        if not os.path.isdir( self.python_dir ):
            os.makedirs( self.python_dir )
        os.system('copy ' + self.source_python + ' ' + self.python_dir + '')
        sys.path.append( self.python_dir )

    def save( self, path='' ):
        if path == '':
            path = os.path.join( self.tmp_dir, self.name + '.study' )
        pfile = open( path, 'w' )
        pickle.dump( self, file=pfile, protocol=0 )
        pfile.close()

    def load( self, path='' ):
        if path == '':
            path = os.path.join( self.tmp_dir, self.name + '.study' )
        pfile = open( path, 'r' )
        self = pickle.load( pfile )
        pfile.close()

    def load_by_name( self, name ):
        self.name = name
        self.rebuild()
        self.load()

    def add_cc_fromDB( self, name ):
        import conecyl
        conecyl = reload( conecyl )
        from conecyl import ConeCyl
        cc = ConeCyl()
        cc.name = name
        cc.fromDB()
        cc.index = len( self.ccs )
        self.ccs.append( cc )
        cc.study = self
        return cc

    def add_cc( self, cc ):
        cc.tmp_dir = self.tmp_dir
        cc.study_dir = self.study_dir
        cc.output_dir = self.output_dir
        cc.study = self
        self.ccs.append( cc )

    def create_models( self, read_outputs = True,
                             last_frame = False,
                             last_cross_section = True ):
        self.rebuild()
        self.rannames = []
        os.chdir( self.output_dir )
        for cc in self.ccs:
            if not cc.check_completed():
                cc.create_model()
                self.rannames.append( cc.jobname )
                cc.write_job( submit = False )
            else:
                cc.create_model()
                if read_outputs:
                    cc.read_outputs( last_frame = last_frame,
                                     last_cross_section = last_cross_section )
        self.create_run_file()
        os.chdir( self.tmp_dir )

    def create_run_file( self ):
        import jobs
        jobs = reload( jobs )
        #
        prefix = os.path.join( self.study_dir, 'run_' + self.name )
        self.run_file_name = prefix + ('_%03d' % self.group) + '.py'
        tmpf = open( self.run_file_name ,'w')
        jobs.print_waitForCompletion( self.study_dir, self.rannames, tmpf )
        tmpf.close()

    def calc_p1( self, step_num = 2 ):
        # 
        import __main__
        self.session = __main__.session
        for i in range(len( self.ccs[0].ploadconf.ploads)):
            pload = self.ccs[0].ploadconf.ploads[i]
            tmp_PLc = []
            for cc in self.ccs:
                cc.read_axial_load_displ( self, step_num )
                pload = cc.ploadconf.ploads[ i ]
                tmp_PLc.append( (pload.pltotal, max(cc.zload)) )
            # estimated P1
            self.p1 = estimate_P1( tmp_PLc )
        return self.p1 

    def create_figures( self ):
        pass

    def plot( self, step_num = 2,
              only_PL        = True,
              cir_criterion  = 'zero_radial_displ',
              cir_threshold  = 0.0125,
              only_cir       = True,
              imp_factor     = False ):
        #
        # circumferential equivalent hole radius
        # 
        import abaqus_functions
        abaqus_functions = reload( abaqus_functions )
        import __main__
        self.session = __main__.session
        numcharts = 5
        if only_PL:
            numcharts = 2
        names = []
        self.xyplots = []
        if not only_PL and not imp_factor:
            for cc in self.ccs:
                for pload in cc.ploadconf.ploads:
                    pload.calc_ehr_and_displ( \
                        cir_criterion = cir_criterion,
                        cir_threshold = cir_threshold,
                        only_cir = only_cir,
                                                )

        for i in range(len( self.ccs[0].ploadconf.ploads)):
            pload = self.ccs[0].ploadconf.ploads[i]
            tmp_cir = []
            tmp_mer = []
            tmp_PLc = []
            tmp_PLc_displ = []
            tmp_dis = []
            xj = []
            yj_cir = []
            yj_mer = []
            ploads = []
            zloads = []
            displs = []
            names.append('Study_%s_%05d_PLoad_%02d_Step_%02d_PL_CURVE' \
                         % ( self.name, self.group, i, step_num ) )
            names.append('Study_%s_%05d_PLoad_%02d_Step_%02d_PL_CURVE_displ' \
                         % ( self.name, self.group, i, step_num ) )
            if not only_PL:
                names.append('Study_%s_%05d_PLoad_%02d_Step_%02d_Equiv_Hole_Radius_cir' \
                             % ( self.name, self.group, i, step_num ) )
                names.append('Study_%s_%05d_PLoad_%02d_Step_%02d_PL_buckle_displ' \
                             % ( self.name, self.group, i, step_num ) )
                if not only_cir:
                    names.append('Study_%s_%05d_PLoad_%02d_Step_%02d_Equiv_Hole_Radius_mer' \
                                 % ( self.name, self.group, i, step_num ) )
            max_ehr_cir = 0.
            max_ehr_mer = 0.
            for cc in self.ccs:
                #
                pload = cc.ploadconf.ploads[ i ]
                step_name = cc.step_name( step_num )
                fb_load = load.find_fb_load( cc.zload )
                #
                if step_num == 2:
                    index = cc.zload.index( fb_load )
                    pload.displ = pload.node.dr[ step_name ][index]
                else:
                    pload.displ = pload.node.dr[ step_name ][-1]
                #
                if not imp_factor:
                    tmp_PLc.append( ( pload.pltotal, 0.001 * fb_load ) )
                else:
                    factor = cc.imperfection_scaling_factors[0]
                    tmp_PLc.append( (        factor, 0.001 * fb_load ) ) 
                #
                tmp_PLc_displ.append( ( abs(pload.displ), 0.001 * fb_load ) ) 
                if not only_PL:
                    ehr_cir = float(pload.ehr_cir)
                    ehr_mer = float(pload.ehr_mer)
                    if ehr_cir > max_ehr_cir:
                        max_ehr_cir = ehr_cir
                    if ehr_mer > max_ehr_mer:
                        max_ehr_mer = ehr_mer
                    xj.append( pload.pltotal )
                    yj_cir.append( ehr_cir )
                    yj_mer.append( ehr_mer )
                    tmp_cir.append( (pload.pltotal, ehr_cir) )
                    tmp_mer.append( (pload.pltotal, ehr_mer) )
                    ploads.append( pload.pltotal )
                    zloads.append( max(cc.zload) )
                    displs.append( pload.displ )
                    tmp_dis.append( (pload.pltotal, pload.displ) )
            # estimated P1
            try:
                est_P1     = estimate_P1( tmp_PLc )
                est_displ  = estimate_P1( tmp_PLc_displ )
            except:
                est_P1 = 0.
                est_displ = 0.
                print 'WARNING - method estimate_P1 failed'
            self.p1 = est_P1
            #
            self.xyplots.append( self.session.XYData(\
                    name = names[ i * numcharts + 0 ],
                    data = tmp_PLc,
                    xValuesLabel = 'Perturbation Load, N',
                    yValuesLabel = 'First Buckling Load, kN',
                    legendLabel  = 'Estimated P1 = %2.2f N' % est_P1 ) )
            self.xyplots.append( self.session.XYData(\
                    name = names[ i * numcharts + 1 ],
                    data = tmp_PLc_displ,
                    xValuesLabel = 'Displacement at PL point (absolute value), mm',
                    yValuesLabel = 'First Buckling Load, kN',
                    legendLabel  = 'Estimated displ = %2.2f mm' % est_displ ) )
            if not only_PL:
                est_P1_cir = estimate_P1( tmp_cir ) 
                self.xyplots.append( self.session.XYData(\
                        name = names[ i * numcharts + 2 ],
                        data = tmp_cir,
                        xValuesLabel = 'Perturbation Load, N',
                        yValuesLabel = 'Equiv Hole Radius - Circumferential, mm',
                        legendLabel  = 'Max Radius = %2.2f mm ///\nEstimated P1 = %2.2f N'\
                                        % (max_ehr_cir, est_P1_cir) ) )
                self.xyplots.append( self.session.XYData(\
                        name = names[ i * numcharts + 3 ],
                        data = tmp_dis,
                        xValuesLabel = 'Perturbation Load, N',
                        yValuesLabel = 'PL_Displacement, mm',
                        legendLabel  = 'Max Displacement  = %2.2f mm' % min(displs) ) )
                if not only_cir:
                    self.xyplots.append( self.session.XYData(\
                            name = names[ i * numcharts + 4 ],
                            data = tmp_mer,
                            xValuesLabel = 'Perturbation Load, N',
                            yValuesLabel = 'Equiv Hole Radius - Meridional, mm',
                            legendLabel  = 'Max Radius = %2.2f mm' % max_ehr_mer ) )

        abaqus_functions.configure_session( session = self.session ) 

